import re
import json
from typing import Dict, List, Set, Optional
from dataclasses import dataclass


@dataclass
class Character:
    """角色信息"""
    name: str
    description: str  # 外观描述
    key_features: List[str]  # 关键特征
    clothing: str  # 服装
    age_group: str  # 年龄段


@dataclass 
class Scene:
    """场景信息"""
    location: str  # 地点
    environment: str  # 环境描述
    lighting: str  # 光线
    atmosphere: str  # 氛围
    time_period: str  # 时间段


@dataclass
class ConsistencyContext:
    """一致性上下文"""
    characters: Dict[str, Character]
    main_scene: Scene
    art_style: str  # 艺术风格
    world_setting: str  # 世界观设定
    color_palette: str  # 色彩基调


class ConsistencyManager:
    """图片一致性管理器"""
    
    def __init__(self):
        self.context: Optional[ConsistencyContext] = None
        self.character_patterns = {
            # 常见角色名称模式
            'names': r'([A-Z][a-z]+|[\u4e00-\u9fff]{1,3})',
            # 外观描述关键词
            'appearance': ['长发', '短发', '黑发', '金发', '棕发', '蓝眼', '棕眼', '高个', '瘦弱', '健壮'],
            # 服装关键词
            'clothing': ['白衣', '黑衣', '制服', '长裙', '短裙', '西装', '便装', '军装'],
            # 年龄关键词  
            'age': ['年轻', '青年', '中年', '少年', '少女', '老者', '孩子']
        }
        self.scene_patterns = {
            # 地点关键词
            'location': ['室内', '室外', '房间', '街道', '森林', '城市', '废墟', '桥上', '山顶'],
            # 环境描述
            'environment': ['昏暗', '明亮', '阳光', '月光', '灯光', '火光', '自然光'],
            # 时间
            'time': ['白天', '夜晚', '黄昏', '清晨', '午后', '深夜'],
            # 天气
            'weather': ['晴朗', '阴天', '雨天', '雪天', '多云', '雾天']
        }
    
    def analyze_chapter_context(self, chapter_content: str) -> ConsistencyContext:
        """分析章节内容，提取一致性上下文"""
        
        # 提取角色信息
        characters = self._extract_characters(chapter_content)
        
        # 提取主要场景
        main_scene = self._extract_main_scene(chapter_content)
        
        # 确定艺术风格
        art_style = "日本吉卜力漫画风格"  # 可以从配置中读取
        
        # 确定世界观设定
        world_setting = self._extract_world_setting(chapter_content)
        
        # 确定色彩基调
        color_palette = self._extract_color_palette(chapter_content)
        
        self.context = ConsistencyContext(
            characters=characters,
            main_scene=main_scene,
            art_style=art_style,
            world_setting=world_setting,
            color_palette=color_palette
        )
        
        return self.context
    
    def _extract_characters(self, text: str) -> Dict[str, Character]:
        """提取角色信息"""
        characters = {}
        
        # 使用正则表达式找到可能的角色名
        name_pattern = r'([A-Z][a-z]+|[\u4e00-\u9fff]{2,3})'
        potential_names = re.findall(name_pattern, text)
        
        # 统计出现频率，筛选主要角色
        name_counts = {}
        for name in potential_names:
            if len(name) >= 2:  # 过滤太短的名字
                name_counts[name] = name_counts.get(name, 0) + 1
        
        # 选择出现频率较高的作为主要角色
        main_characters = [name for name, count in name_counts.items() if count >= 2]
        
        for char_name in main_characters[:3]:  # 最多保留3个主要角色
            # 查找角色相关的描述
            char_descriptions = self._find_character_descriptions(text, char_name)
            
            characters[char_name] = Character(
                name=char_name,
                description=char_descriptions.get('appearance', '年轻人'),
                key_features=char_descriptions.get('features', []),
                clothing=char_descriptions.get('clothing', '普通服装'),
                age_group=char_descriptions.get('age', '年轻')
            )
        
        return characters
    
    def _find_character_descriptions(self, text: str, char_name: str) -> Dict[str, any]:
        """查找角色相关描述"""
        descriptions = {
            'appearance': '',
            'features': [],
            'clothing': '',
            'age': '年轻'
        }
        
        # 查找包含角色名的句子
        sentences = re.split(r'[。！？]', text)
        char_sentences = [s for s in sentences if char_name in s]
        
        for sentence in char_sentences:
            # 查找外观描述
            for keyword in self.character_patterns['appearance']:
                if keyword in sentence:
                    descriptions['features'].append(keyword)
            
            # 查找服装描述
            for keyword in self.character_patterns['clothing']:
                if keyword in sentence:
                    descriptions['clothing'] = keyword
                    break
            
            # 查找年龄描述
            for keyword in self.character_patterns['age']:
                if keyword in sentence:
                    descriptions['age'] = keyword
                    break
        
        # 生成综合外观描述
        if descriptions['features']:
            descriptions['appearance'] = f"{descriptions['age']}{char_name}，{','.join(descriptions['features'][:2])}"
        else:
            descriptions['appearance'] = f"{descriptions['age']}{char_name}"
        
        return descriptions
    
    def _extract_main_scene(self, text: str) -> Scene:
        """提取主要场景信息"""
        
        # 查找地点信息
        location = '室内'
        for keyword in self.scene_patterns['location']:
            if keyword in text:
                location = keyword
                break
        
        # 查找环境描述
        environment = '普通环境'
        lighting = '自然光'
        for keyword in self.scene_patterns['environment']:
            if keyword in text:
                lighting = keyword
                break
        
        # 查找时间信息
        time_period = '白天'
        for keyword in self.scene_patterns['time']:
            if keyword in text:
                time_period = keyword
                break
        
        # 确定氛围
        atmosphere = '平静'
        if any(word in text for word in ['紧张', '危险', '对峙', '威胁']):
            atmosphere = '紧张'
        elif any(word in text for word in ['温暖', '平和', '宁静']):
            atmosphere = '温暖'
        elif any(word in text for word in ['神秘', '诡异', '阴暗']):
            atmosphere = '神秘'
        
        return Scene(
            location=location,
            environment=environment,
            lighting=lighting,
            atmosphere=atmosphere,
            time_period=time_period
        )
    
    def _extract_world_setting(self, text: str) -> str:
        """提取世界观设定"""
        if any(word in text for word in ['废土', '末世', '废墟', '资源匮乏']):
            return '末世废土世界'
        elif any(word in text for word in ['现代', '都市', '现在']):
            return '现代都市'
        elif any(word in text for word in ['古代', '古老', '传统']):
            return '古代世界'
        else:
            return '现代世界'
    
    def _extract_color_palette(self, text: str) -> str:
        """提取色彩基调"""
        if any(word in text for word in ['昏暗', '阴暗', '黑暗']):
            return '低饱和度，偏暗色调'
        elif any(word in text for word in ['明亮', '阳光', '温暖']):
            return '高饱和度，暖色调'
        elif any(word in text for word in ['冷', '蓝', '月光']):
            return '冷色调，蓝绿为主'
        else:
            return '自然色调'
    
    def generate_consistent_prompt(self, paragraph: str, paragraph_index: int) -> str:
        """生成具有一致性的图片提示词"""
        if not self.context:
            raise ValueError("需要先调用 analyze_chapter_context 来设置上下文")
        
        # 分析当前段落的具体内容
        paragraph_analysis = self._analyze_paragraph(paragraph)
        
        # 提取段落的核心视觉描述
        core_visual_description = self._extract_visual_description(paragraph)
        
        # 构建完整提示词
        # 1. 艺术风格
        prompt_parts = [self.context.art_style]
        
        # 2. 具体的视觉场景描述
        if core_visual_description:
            prompt_parts.append(core_visual_description)
        
        # 3. 角色信息（如果存在）
        if paragraph_analysis['characters']:
            for char_name in paragraph_analysis['characters']:
                if char_name in self.context.characters:
                    char = self.context.characters[char_name]
                    prompt_parts.append(f"{char.description}")
        
        # 4. 场景设定和氛围
        scene_info = self._get_scene_info(paragraph, paragraph_analysis)
        prompt_parts.append(scene_info)
        
        # 5. 具体物品和动作（如果有的话）
        objects_and_actions = self._extract_objects_and_actions(paragraph)
        if objects_and_actions:
            prompt_parts.append(objects_and_actions)
        
        # 6. 技术质量要求
        prompt_parts.extend([
            self.context.color_palette,
            "电影级画质",
            "细腻光影效果",
            "高质量渲染"
        ])
        
        return "，".join(prompt_parts)
    
    def _analyze_paragraph(self, paragraph: str) -> Dict[str, any]:
        """分析段落内容"""
        analysis = {
            'characters': [],
            'actions': [],
            'objects': [],
            'emotions': []
        }
        
        # 查找角色
        for char_name in self.context.characters.keys():
            if char_name in paragraph:
                analysis['characters'].append(char_name)
        
        # 查找动作
        action_keywords = ['站着', '坐着', '走', '跑', '拿', '举', '看', '听', '说', '喝', '吃', '抱']
        for action in action_keywords:
            if action in paragraph:
                analysis['actions'].append(action)
        
        # 查找物品
        object_keywords = ['可乐', '刀', '手环', '背包', '罐头', '武器', '饮料']
        for obj in object_keywords:
            if obj in paragraph:
                analysis['objects'].append(obj)
        
        return analysis
    
    def _get_scene_info(self, paragraph: str, analysis: Dict) -> str:
        """获取场景信息"""
        # 如果段落中有特定场景描述，使用段落的；否则使用全局场景
        scene_desc = self.context.main_scene.location
        
        # 检查段落中是否有特定地点描述
        for location in self.scene_patterns['location']:
            if location in paragraph:
                scene_desc = location
                break
        
        # 添加光线和氛围
        lighting = self.context.main_scene.lighting
        for light in self.scene_patterns['environment']:
            if light in paragraph:
                lighting = light
                break
        
        return f"{scene_desc}，{lighting}，{self.context.main_scene.atmosphere}氛围"
    
    def _extract_visual_description(self, paragraph: str) -> str:
        """提取段落的核心视觉描述"""
        visual_elements = []
        
        # 查找具体的视觉场景描述
        if '刀锋抵在' in paragraph and '额头' in paragraph:
            visual_elements.append("冰冷刀锋抵在额头的紧张对峙场景")
        
        if '可乐' in paragraph and ('喝' in paragraph or '饮' in paragraph):
            visual_elements.append("人物饮用可乐的动作特写")
        
        if '眼神' in paragraph and ('锐利' in paragraph or '凌厉' in paragraph):
            visual_elements.append("锐利眼神特写")
        
        if '手环' in paragraph and '蓝光' in paragraph:
            visual_elements.append("手腕上闪烁蓝光的神秘手环")
        
        if '罐头' in paragraph and '掏出' in paragraph:
            visual_elements.append("从背包掏出食物罐头的动作")
        
        if '微笑' in paragraph:
            visual_elements.append("人物露出微笑的表情")
        
        # 查找情绪描述
        if '紧张' in paragraph or '威胁' in paragraph:
            visual_elements.append("紧张威胁的氛围")
        elif '满足' in paragraph:
            visual_elements.append("满足的表情")
        elif '警惕' in paragraph:
            visual_elements.append("警惕戒备的神态")
        
        # 查找环境描述
        if '昏暗' in paragraph:
            visual_elements.append("昏暗的室内环境")
        elif '光线' in paragraph:
            visual_elements.append("戏剧性的光影效果")
        
        return "，".join(visual_elements) if visual_elements else "日常生活场景"
    
    def _extract_objects_and_actions(self, paragraph: str) -> str:
        """提取具体的物品和动作描述"""
        elements = []
        
        # 重要物品
        if '可乐' in paragraph:
            elements.append("可乐罐")
        if '刀' in paragraph:
            elements.append("锋利的刀具")
        if '手环' in paragraph:
            elements.append("科技感手环")
        if '罐头' in paragraph:
            elements.append("食物罐头")
        if '背包' in paragraph:
            elements.append("背包")
        
        # 具体动作
        if '抵在' in paragraph:
            elements.append("威胁动作")
        elif '喝' in paragraph or '饮' in paragraph:
            elements.append("饮用动作")
        elif '掏出' in paragraph:
            elements.append("拿取动作")
        elif '放下' in paragraph:
            elements.append("放下动作")
        
        return "，".join(elements) if elements else None
    
    def _extract_action_description(self, paragraph: str) -> str:
        """提取动作描述（保留兼容性）"""
        # 这个方法现在主要用于场景分类
        if '对话' in paragraph or '"' in paragraph:
            return "对话场景"
        elif '抵在' in paragraph:
            return "紧张对峙场景"
        elif '喝' in paragraph or '饮' in paragraph:
            return "饮用场景"
        elif '掏出' in paragraph:
            return "拿取物品场景"
        else:
            return "日常场景"
    
    def save_context(self, output_dir: str, chapter_hash: str):
        """保存一致性上下文到文件"""
        if not self.context:
            return
        
        context_file = f"{output_dir}/consistency_context_{chapter_hash}.json"
        
        # 转换为可序列化的字典
        context_dict = {
            'characters': {
                name: {
                    'name': char.name,
                    'description': char.description,
                    'key_features': char.key_features,
                    'clothing': char.clothing,
                    'age_group': char.age_group
                } for name, char in self.context.characters.items()
            },
            'main_scene': {
                'location': self.context.main_scene.location,
                'environment': self.context.main_scene.environment,
                'lighting': self.context.main_scene.lighting,
                'atmosphere': self.context.main_scene.atmosphere,
                'time_period': self.context.main_scene.time_period
            },
            'art_style': self.context.art_style,
            'world_setting': self.context.world_setting,
            'color_palette': self.context.color_palette
        }
        
        with open(context_file, 'w', encoding='utf-8') as f:
            json.dump(context_dict, f, ensure_ascii=False, indent=2)
        
        print(f"📋 一致性上下文已保存: {context_file}")